from argparse import Namespace

import torch
import torchaudio
import torchaudio.compliance.kaldi as kaldi
import yaml
from wenet.utils.init_model import init_model
from wenet.utils.init_tokenizer import init_tokenizer

config = 'wenet_fired/train.yaml'
args = Namespace(checkpoint='wenet_fired/wenet_firered.pt')

with open(config, 'r') as fin:
    configs = yaml.load(fin, Loader=yaml.FullLoader)

model, _ = init_model(args, configs)

model.eval()
audio_file = '../../test/resources/aishell-BAC009S0724W0121.wav'
waveform, sample_rate = torchaudio.load(audio_file, normalize=False)
waveform = waveform.to(torch.float)
feats = kaldi.fbank(waveform,
                    num_mel_bins=80,
                    frame_length=25,
                    frame_shift=10,
                    energy_floor=0.0,
                    sample_frequency=16000)

feats = feats.unsqueeze(0)
feats_lens = torch.tensor([feats.shape[1]], dtype=torch.int64)
results = model.decode(methods=['attention'],
                       speech=feats,
                       speech_lengths=feats_lens,
                       beam_size=10)

tokenizer = init_tokenizer(configs)

for mode, hyps in results.items():
    print('{} {}'.format(mode, tokenizer.detokenize(hyps[0].tokens)))

from fireredasr.models.fireredasr import FireRedAsr

wav_paths = ['../../test/resources/aishell-BAC009S0724W0121.wav']

red_model = FireRedAsr.from_pretrained("aed", "pretrained_models/FireRedASR-AED-L")
red_feats, lengths, durs = red_model.feat_extractor(wav_paths)
red_enc_outputs,  _, enc_mask = model.model.encoder(red_feats, lengths)


wenet_encoder = model.encoder
wenet_enc_outs, _ = wenet_encoder(feats,  feats_lens)

torch.allclose(wenet_enc_outs, red_enc_outputs) # True